import subprocess
import sys
import platform
import re
from typing import Dict, List, Union

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
    from mcp.server.fastmcp import FastMCP
    MCP_AVAILABLE = True
except ImportError:
    print("警告: 未安装MCP库。请运行 'pip install mcp' 安装。")
    print("在此期间，您仍然可以使用核心网络检测功能。")
    ClientSession = None
    StdioServerParameters = None
    stdio_client = None
    FastMCP = None
    MCP_AVAILABLE = False

class NetworkConnectivityChecker:
    """
    网络连通性检测工具类
    """
    
    def ping_host(self, host: str, count: int = 1, timeout: int = 3) -> Dict[str, Union[bool, float, str]]:
        """
        检测主机连通性
        
        Args:
            host: 目标主机IP或域名
            count: ping的次数
            timeout: 超时时间(秒)
            
        Returns:
            包含检测结果的字典
        """
        try:
            # 判断操作系统
            system = platform.system().lower()
            
            if system == "windows":
                # Windows系统ping命令
                cmd = ["ping", "-n", str(count), "-w", str(timeout * 1000), host]
            else:
                # Unix/Linux/Mac系统ping命令
                cmd = ["ping", "-c", str(count), "-W", str(timeout), host]
            
            # 执行ping命令
            output = subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True)
            
            # 解析结果
            result = self._parse_ping_output(output, system)
            result["success"] = True
            result["host"] = host
            
            return result
            
        except subprocess.CalledProcessError as e:
            # ping失败
            return {
                "success": False,
                "host": host,
                "error": f"Ping failed with return code: {e.returncode}",
                "output": e.output if e.output else ""
            }
        except Exception as e:
            # 其他异常
            return {
                "success": False,
                "host": host,
                "error": str(e)
            }
    
    def _parse_ping_output(self, output: str, system: str) -> Dict[str, Union[float, str]]:
        """
        解析ping命令输出
        
        Args:
            output: ping命令输出
            system: 操作系统类型
            
        Returns:
            解析结果字典
        """
        result = {}
        
        if system == "windows":
            # 解析Windows的ping输出
            time_match = re.search(r"时间[=<](\d+)ms", output)
            if time_match:
                result["response_time"] = float(time_match.group(1))
            
            loss_match = re.search(r"丢失 = (\d+)", output)
            if loss_match:
                result["packet_loss"] = int(loss_match.group(1))
        else:
            # 解析Unix/Linux/Mac的ping输出
            time_match = re.search(r"time=(\d+\.?\d*)\s*ms", output)
            if time_match:
                result["response_time"] = float(time_match.group(1))
            
            loss_match = re.search(r"(\d+)% packet loss", output)
            if loss_match:
                result["packet_loss"] = int(loss_match.group(1))
        
        return result
    
    def check_port(self, host: str, port: int, timeout: int = 3) -> Dict[str, Union[bool, str]]:
        """
        检查主机端口是否开放
        
        Args:
            host: 目标主机IP或域名
            port: 端口号
            timeout: 超时时间(秒)
            
        Returns:
            包含检测结果的字典
        """
        import socket
        
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            result = sock.connect_ex((host, port))
            sock.close()
            
            if result == 0:
                return {
                    "success": True,
                    "host": host,
                    "port": port,
                    "open": True,
                    "message": f"Port {port} is open"
                }
            else:
                return {
                    "success": True,
                    "host": host,
                    "port": port,
                    "open": False,
                    "message": f"Port {port} is closed"
                }
        except Exception as e:
            return {
                "success": False,
                "host": host,
                "port": port,
                "error": str(e)
            }

# 如果安装了MCP库，则提供MCP工具封装
if MCP_AVAILABLE:
    # 创建FastMCP实例
    mcp = FastMCP("Network Connectivity Checker")
    
    # 使用装饰器将函数注册为MCP工具
    @mcp.tool()
    def ping_host(host: str, count: int = 1, timeout: int = 3) -> Dict[str, Union[bool, float, str]]:
        """
        检查主机连通性 (MCP工具版本)
        
        Args:
            host: 目标主机IP或域名
            count: ping的次数
            timeout: 超时时间(秒)
            
        Returns:
            包含检测结果的字典
        """
        checker = NetworkConnectivityChecker()
        return checker.ping_host(host, count, timeout)
    
    @mcp.tool()
    def check_port_connectivity(host: str, port: int, timeout: int = 3) -> Dict[str, Union[bool, str]]:
        """
        检查主机端口是否开放 (MCP工具版本)
        
        Args:
            host: 目标主机IP或域名
            port: 端口号
            timeout: 超时时间(秒)
            
        Returns:
            包含检测结果的字典
        """
        checker = NetworkConnectivityChecker()
        return checker.check_port(host, port, timeout)

def main():
    """
    主函数，用于命令行测试
    """
    checker = NetworkConnectivityChecker()
    
    if len(sys.argv) < 2:
        print("使用方法:")
        print("  python check_port.py <host> [count] [timeout]")
        print("  python check_port.py <host> port <port_number> [timeout]")
        print("\n示例:")
        print("  python check_port.py 8.8.8.8")
        print("  python check_port.py www.google.com 4 5")
        print("  python check_port.py 8.8.8.8 port 53")
        return
    
    host = sys.argv[1]
    
    if len(sys.argv) >= 3 and sys.argv[2] == "port":
        # 端口检测模式
        if len(sys.argv) >= 4:
            port = int(sys.argv[3])
            timeout = int(sys.argv[4]) if len(sys.argv) > 4 else 3
            result = checker.check_port(host, port, timeout)
        else:
            print("端口检测需要指定端口号")
            return
    else:
        # Ping检测模式
        count = int(sys.argv[2]) if len(sys.argv) > 2 else 1
        timeout = int(sys.argv[3]) if len(sys.argv) > 3 else 3
        result = checker.ping_host(host, count, timeout)
    
    # 打印结果
    print(result)

# 如果安装了MCP库，添加启动服务器的代码
if MCP_AVAILABLE:
    @mcp.tool()
    def start_mcp_server():
        """
        启动MCP服务器
        """
        print("启动MCP服务器...")
        mcp.run()

if __name__ == "__main__":
    main()